# File: src/mcp_client/session_manager.py

import asyncio
import logging
import os
import sys
from typing import Optional, Any
import httpx
import anyio
from anyio.abc import TaskGroup # <--- 修改导入方式

# --- 核心导入 ---
from ..core.config import load_config, get_config_value
from ..core.exceptions import MCPClientError, ConfigError
from ..core.logging_setup import get_logger

# --- MCP SDK 导入 ---
try:
    from mcp import ClientSession, types as mcp_types
    from mcp.client.sse import sse_client
    from mcp.shared.exceptions import McpError
except ImportError as e:
    logging.basicConfig(level=logging.CRITICAL)
    logging.critical(f"Failed to import MCP client components (sse): {e}. Ensure SDK installed.", exc_info=True)
    raise ImportError("MCP SDK client components (sse) not found.") from e

logger = get_logger("mcp_client")

# --- 全局客户端状态 ---
_active_session: Optional[ClientSession] = None
_connection_ctx_manager: Optional[Any] = None
# *** 修改类型提示 ***
_session_task_group: Optional[TaskGroup] = None # 使用导入的 TaskGroup
_client_lock = asyncio.Lock()

async def initialize_mcp_client() -> ClientSession:
    """初始化 MCP 客户端连接 (SSE 模式) 并启动其内部任务，然后返回活动会话"""
    global _active_session, _connection_ctx_manager, _session_task_group
    async with _client_lock:
        if _active_session:
            logger.debug("MCP client session (SSE) already exists.")
            return _active_session

        logger.info("Initializing MCP Client Connection (SSE Mode)...")
        await _close_internal_resources() # 先关闭旧资源

        session: Optional[ClientSession] = None

        try:
            # --- 1. 启动 SSE Transport ---
            logger.debug("Forcing config reload before getting SSE URL...")
            config = load_config(force_reload=True)
            base_url = config.get("mcp_client", {}).get("sse", {}).get("base_url")
            logger.info(f"Read SSE base_url from config: '{base_url}'")
            if not base_url:
                raise ConfigError("SSE connection requires mcp_client.sse.base_url to be configured.")

            logger.debug("Entering sse_client context manager...")
            temp_ctx_manager = sse_client(url=base_url)
            read_stream, write_stream = await temp_ctx_manager.__aenter__()
            _connection_ctx_manager = temp_ctx_manager
            logger.debug("SSE streams obtained via sse_client context.")

            # --- 2. 创建 ClientSession ---
            session = ClientSession(read_stream, write_stream)
            logger.debug("ClientSession instance created.")

            # --- 3. 手动启动 ClientSession 后台任务 ---
            try:
                logger.debug("Manually starting ClientSession TaskGroup and _receive_loop...")
                tg = anyio.create_task_group()
                await tg.__aenter__()
                _session_task_group = tg # 赋值
                # *** 在使用前添加断言 ***
                assert _session_task_group is not None, "Task group should be initialized here"
                _session_task_group.start_soon(session._receive_loop) # 现在 Pyright 不会报错
                logger.debug("ClientSession _receive_loop started in TaskGroup.")
            except Exception as tg_err:
                logger.error("Failed to manually start ClientSession TaskGroup/receive_loop", exc_info=True)
                if _connection_ctx_manager:
                    try: await _connection_ctx_manager.__aexit__(None, None, None)
                    except Exception: pass
                _connection_ctx_manager = None
                raise MCPClientError("Failed to start internal session tasks") from tg_err

            # --- 4. 执行 MCP 初始化握手 ---
            logger.info("Initializing MCP session handshake (SSE)...")
            try:
                init_result = await session.initialize()
                logger.info(f"MCP session handshake successful. Server: {init_result.serverInfo.name} v{init_result.serverInfo.version}")
            except Exception as init_err:
                logger.error(f"Error during session.initialize(): {init_err}", exc_info=True)
                await _close_internal_resources()
                raise MCPClientError(f"MCP handshake failed: {init_err}") from init_err

            # --- 5. 保存活动会话 ---
            _active_session = session
            logger.info(f"MCP Client Session (SSE) initialized successfully.")
            return _active_session

        except Exception as e:
            logger.critical(f"Failed to initialize MCP client session (SSE): {type(e).__name__} - {e}", exc_info=True)
            await _close_internal_resources()
            if isinstance(e, (MCPClientError, ConfigError)):
                 raise e
            else:
                 raise MCPClientError(f"SSE initialization failed unexpectedly: {e}") from e


async def _close_internal_resources():
    """内部清理函数，安全地关闭 session task group 和 sse transport context"""
    global _active_session, _connection_ctx_manager, _session_task_group
    logger.debug("Entering _close_internal_resources...")
    session_to_clear = _active_session
    ctx_manager_to_exit = _connection_ctx_manager
    tg_to_exit = _session_task_group
    _active_session = None
    _connection_ctx_manager = None
    _session_task_group = None

    if tg_to_exit:
        logger.info("Exiting ClientSession internal TaskGroup (_receive_loop)...")
        try:
            tg_to_exit.cancel_scope.cancel()
            await tg_to_exit.__aexit__(None, None, None)
            logger.info("ClientSession TaskGroup exited.")
        except Exception as tg_exit_err:
            logger.error(f"Error exiting ClientSession TaskGroup: {tg_exit_err}", exc_info=True)

    if ctx_manager_to_exit:
        logger.info("Exiting MCP client connection context manager (SSE)...")
        try:
            await ctx_manager_to_exit.__aexit__(None, None, None)
            logger.info("Connection context manager (SSE) exited successfully.")
        except Exception as e:
            logger.error(f"Error exiting MCP client connection context (SSE): {e}", exc_info=True)

    if session_to_clear:
        logger.debug(f"Cleared active session reference for session {id(session_to_clear)}.")
    logger.debug("Finished _close_internal_resources.")


async def close_mcp_client() -> None:
    """关闭活动的 MCP 客户端会话和 SSE 连接上下文"""
    async with _client_lock:
        if _active_session or _connection_ctx_manager or _session_task_group:
            logger.info("Closing MCP client resources (SSE)...")
            await _close_internal_resources()
            logger.info("Finished closing MCP client resources (SSE).")
        else:
            logger.debug("MCP client close called, but no active resources found.")


async def get_mcp_session() -> ClientSession:
     """返回活动的 MCP 客户端会话。如果需要则异步初始化。"""
     async with _client_lock:
         if _active_session:
             return _active_session
     return await initialize_mcp_client()